#include <iostream>
#include <ceres/ceres.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <glog/logging.h>

using namespace pybind11::literals;
namespace py = pybind11;

class CustomFunction : public ceres::FirstOrderFunction {
    public:
        CustomFunction(py::function func,
                       py::function grad,
                       unsigned long num_params)
                        : function_(func),
                          gradient_(grad),
                          num_params_(num_params) {}
        virtual ~CustomFunction() override {}
        virtual bool Evaluate(const double* parameters,
                              double* cost,
                              double* gradient) const override {
            auto arr = py::array(py::buffer_info(
                const_cast<double*>(parameters),
                sizeof(double),
                py::format_descriptor<double>::value,
                1,
                { num_params_ },
                { sizeof(double) }
            ));

            auto out = function_(arr);

            // If the cost evaluation fails, the function should return None.
            // Returning false means that Ceres should try other parameters.
            if (out == py::none()) {
                return false;
            }

            cost[0] = out.cast<double>();

            // If gradient is not null, we are supposed to
            // calculate the gradient.
            if (gradient) {
                auto out = gradient_(arr);

                // We also check the gradient for None.
                if (out == py::none()) {
                    return false;
                }

                auto arr = out.cast<py::array_t<double>>();
                auto buf = static_cast<double*>(arr.request().ptr);
                for (unsigned long i = 0; i < num_params_; i++) {
                    gradient[i] = buf[i];
                }
            }
            return true;
        }
        virtual int NumParameters() const override {
            return num_params_;
        }
    private:
        py::function function_;
        py::function gradient_;
        unsigned long num_params_;
};

py::object optimize(py::function func, py::function grad, py::buffer x0) {
    auto buf1 = x0.request();
    if (buf1.ndim != 1) {
        throw std::runtime_error("Number of dimensions of x0 must be one");
    }
    ceres::GradientProblemSolver::Options options;
    options.minimizer_progress_to_stdout = true;
    ceres::GradientProblemSolver::Summary summary;
    ceres::GradientProblem problem(new CustomFunction(func, grad, buf1.size));

    auto result = py::array_t<double>(buf1.size);
    auto buf2 = result.request();

    double* data = static_cast<double*>(buf2.ptr);
    double* inputs = static_cast<double*>(buf1.ptr);

    for (unsigned long i = 0; i < buf1.size; i++) {
        data[i] = inputs[i];
    }

    ceres::Solve(options, problem, data, &summary);

    auto iterations = summary.iterations;
    int nfev = 0;
    int ngev = 0;
    for (auto& summ: iterations) {
        nfev += summ.line_search_function_evaluations;
        ngev += summ.line_search_gradient_evaluations;
    }

    auto OptimizeResult = py::module::import("scipy.optimize").attr("OptimizeResult");

    py::dict out("x"_a = result,
                 "success"_a = summary.termination_type == ceres::CONVERGENCE ||
                               summary.termination_type == ceres::USER_SUCCESS,
                 "status"_a = (int)summary.termination_type,
                 "message"_a = summary.message,
                 "fun"_a = summary.final_cost,
                 "nfev"_a = nfev,
                 "ngev"_a = ngev);

    return OptimizeResult(out);
}


//input:
// vector 3d点坐标
// vector 2d点坐标
// 旋转矩阵 r
// 平移向量 t
py::object optimize_ceres(py::function func,py::buffer RR,py::buffer tt, py::array_t<double> obj,py::array_t<double> uv,py::array_t<int> num)
{  
   double RR_input[3]={0,0,0};//rpy
   double tt_input[3]={0,0,0};//tt

   auto num_temp= num.request();
   int* num_input_temp = static_cast<int*>(num_temp.ptr);
    
   double obj_input[*num_input_temp][3];//三维坐标点
   double uv_input[*num_input_temp][2];//二维坐标点

/*
   auto buf1 = RR.request();
    if (buf1.ndim != 1) {
        throw std::runtime_error("Number of dimensions of x0 must be one");
    }
   double* inputs = static_cast<double*>(buf1.ptr);
   for (unsigned long i = 0; i < buf1.size; i++) {
        data[i] = inputs[i];
    }
*/

   auto RR_temp = RR.request(); 
   auto tt_temp = tt.request();
   auto obj_temp = obj.request();
   auto uv_temp = uv.request();
   
   if (RR_temp.ndim != 1 ) {
        throw std::runtime_error("Number of dimensions of x0 must be one");
    }


   double* RR_input_temp = static_cast<double*>(RR_temp.ptr);
   double* tt_input_temp = static_cast<double*>(tt_temp.ptr);
   double* obj_input_temp = static_cast<double*>(obj_temp.ptr);
   double* uv_input_temp = static_cast<double*>(uv_temp.ptr);
   


   for (unsigned long i = 0; i < RR_temp.size; i++) {
        RR_input[i] = RR_input_temp[i];
        tt_input[i] = tt_input_temp[i];
    }
    std::cout<<"dim:"<<obj_temp.ndim<<std::endl;


    for (unsigned long i = 0; i < obj_temp.ndim; i++)
    {
      obj_input[i][0]=obj_input_temp[3*i];
      obj_input[i][1]=obj_input_temp[3*i+1];
      obj_input[i][2]=obj_input_temp[3*i+2];
      uv_input[i][0]=uv_input_temp[3*i];
      uv_input[i][1]=uv_input_temp[3*i+1];
    }



    ceres::Solver::Options d_opt;
    d_opt.linear_solver_type=ceres::DENSE_QR;
    d_opt.minimizer_progress_to_stdout=true;
    ceres::Solver::Summary sum;
    std::cout<<"开始求解"<<std::endl;
    auto result = py::array_t<double>(11);
    std::cout<<"求解出结果"<<std::endl;
    auto OptimizeResult = py::module::import("scipy.optimize").attr("OptimizeResult");
    py::dict out("x"_a = result);

    return OptimizeResult(out);

}




namespace py = pybind11;

PYBIND11_PLUGIN(ceres) {
    py::module m("ceres", "Python bindings to the Ceres-Solver minimizer.");
    google::InitGoogleLogging("ceres");

    m.def("optimize", &optimize, "Optimizes the function");
    m.def("optimize_ceres", &optimize_ceres, "Optimizes the function");
    return m.ptr();
}

